
import argparse
import os
import json
import cv2
from tqdm import tqdm
# import pycocotools.mask as maskUtils
import numpy as np
import cv2
import pandas as pd
import shutil


def main():
    parser = argparse.ArgumentParser(description="Convert OpenImages annotation to spire annotation")
    parser.add_argument(
        "--openimages-imgdir",
        default="C:/dataset/BB-OpenImages-train-v240218",
        help="path to openimages image dir",
        # required=True
    )
    parser.add_argument(
        "--openimages-csv",
        default="C:/dataset/boxes/train-annotations-bbox.csv",
        help="path to openimages csv",
    )
    parser.add_argument(
        "--output-dir",
        default="C:/dataset/BB-OpenImages-train-v240218",
        help="path to spire home dir",
    )
    args = parser.parse_args()

    """
    output_img_dir = os.path.join(args.output_dir, 'scaled_images')
    output_ann_dir = os.path.join(args.output_dir, 'annotations')
    if os.path.exists(output_img_dir):
        shutil.rmtree(output_img_dir)
    if os.path.exists(output_ann_dir):
        shutil.rmtree(output_ann_dir)

    if not os.path.exists(output_img_dir):
        os.makedirs(output_img_dir)
    if not os.path.exists(output_ann_dir):
        os.makedirs(output_ann_dir)
    """

    images_w_annos = dict()
    # f = open(args.openimages_csv, 'r')
    df = pd.read_csv(args.openimages_csv)
    for (ImageID, Source, LabelName, Confidence, XMin, XMax, YMin, YMax,
        IsOccluded, IsTruncated, IsGroupOf, IsDepiction, IsInside) \
            in zip(df.ImageID, df.Source, df.LabelName, df.Confidence,
                df.XMin, df.XMax, df.YMin, df.YMax, df.IsOccluded, df.IsTruncated,
                df.IsGroupOf, df.IsDepiction, df.IsInside):
        # print(ImageID, Source, LabelName, Confidence, XMin, XMax, YMin, YMax)
        if ImageID in images_w_annos.keys():
            images_w_annos[ImageID].append([LabelName, Confidence, XMin, XMax, YMin, YMax,
                                            IsOccluded, IsTruncated, IsGroupOf, IsDepiction, IsInside])
        else:
            images_w_annos[ImageID] = [[LabelName, Confidence, XMin, XMax, YMin, YMax,
                                        IsOccluded, IsTruncated, IsGroupOf, IsDepiction, IsInside]]

    sub_dirs = [
        'train_0',
        'train_1',
        'train_2',
        'train_3',
        'train_4',
        'train_5',
        'train_6',
        'train_7',
        'train_8',
        'train_9',
        'train_a',
        'train_b',
        'train_c',
        'train_d',
        'train_e',
        'train_f'
    ]
    for key in tqdm(images_w_annos.keys()):
        # print(key)
        sub_dir = ''
        for sd in sub_dirs:
            img_fn = os.path.join(args.openimages_imgdir, sd, 'scaled_images', key + '.jpg')
            if os.path.exists(img_fn):
                sub_dir = sd
                break
        assert len(sub_dir) > 0
        img = cv2.imread(img_fn)
        spire_dict = {}
        spire_dict['file_name'] = key + '.jpg'
        h, w = img.shape[0], img.shape[1]
        spire_dict['height'], spire_dict['width'] = img.shape[0], img.shape[1]
        spire_dict['annos'] = []
        for anno in images_w_annos[key]:
            (LabelName, Confidence, XMin, XMax, YMin, YMax,
             IsOccluded, IsTruncated, IsGroupOf, IsDepiction, IsInside) = anno
            spire_anno = {}
            ox = XMin * w
            oy = YMin * h
            ow = (XMax - XMin) * w
            oh = (YMax - YMin) * h
            spire_anno['area'] = ow * oh
            spire_anno['bbox'] = [ox, oy, ow, oh]

            spire_anno['is_occluded'] = IsOccluded
            spire_anno['is_truncated'] = IsTruncated
            spire_anno['is_group_of'] = IsGroupOf
            spire_anno['is_depiction'] = IsDepiction
            spire_anno['is_inside'] = IsInside
            spire_anno['score'] = Confidence
            spire_anno['category_name'] = LabelName
            spire_dict['annos'].append(spire_anno)

        output_ann_dir = str(os.path.join(args.output_dir, sub_dir, 'annotations'))
        if not os.path.exists(output_ann_dir):
            os.makedirs(output_ann_dir)
        output_fn = os.path.join(output_ann_dir, key + '.jpg' + '.json')
        with open(output_fn, "w") as f:
            json.dump(spire_dict, f)

        # img_fn = os.path.join(output_img_dir, key + '.jpg')
        # cv2.imwrite(img_fn, img)


if __name__ == '__main__':
    main()
